# mobile_net_v2_codes.py
import os
import glob
import time
from typing import List          # ★ 让 IDE 识别列表类型

from PIL import Image
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, models
from tqdm import tqdm
import matplotlib.pyplot as plt

# ---------- 路径/超参 ----------
TRAIN_DIR = r'D:\codes\learntorch\datas128\seg_train'
VAL_DIR = r'D:\codes\learntorch\datas128\seg_val'
PRED_DIR = r'D:\codes\learntorch\datas128\seg_pred'
SAVE_DIR = r'D:\codes\learntorch\printer'
MODEL_PATH = os.path.join(SAVE_DIR, 'best_mobilenetv2.pth')
IMG_SIZE, BATCH_SIZE, NUM_EPOCHS, LR, PATIENCE = 128, 32, 20, 1e-3, 5

if not torch.cuda.is_available():
    raise RuntimeError('❌ 未检测到 GPU')
DEVICE = torch.device('cuda')
os.makedirs(SAVE_DIR, exist_ok=True)

class FolderDataset(Dataset):
    def __init__(self, root, transform=None):
        self.transform = transform
        self.classes = sorted([d for d in os.listdir(root) if os.path.isdir(os.path.join(root, d))])
        self.class2idx = {c: i for i, c in enumerate(self.classes)}
        self.samples: List[str] = []                       # ★ 先建空列表
        for cls in self.classes:
            # ★ 用 extend 追加，不要重新赋值
            self.samples.extend(list(glob.glob(os.path.join(root, cls, '*'))))

    def __len__(self) -> int:
        return len(self.samples)

    def __getitem__(self, idx):
        path = self.samples[idx]
        img = Image.open(path).convert('RGB')
        cls_name = os.path.basename(os.path.dirname(path))
        label = self.class2idx[cls_name]
        if self.transform:
            img = self.transform(img)
        return img, label


# ---------- 数据增强 ----------
train_tf = transforms.Compose([
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
val_tf = transforms.Compose([
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

train_set = FolderDataset(TRAIN_DIR, train_tf)
val_set = FolderDataset(VAL_DIR, val_tf)
train_loader = DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True, num_workers=0, pin_memory=True)
val_loader = DataLoader(val_set, batch_size=BATCH_SIZE, shuffle=False, num_workers=0, pin_memory=True)
NUM_CLASSES = len(train_set.classes)

# ---------- 模型 ----------
from torchvision.models import MobileNet_V2_Weights


def build_model(pretrained: bool = True):
    weights = MobileNet_V2_Weights.IMAGENET1K_V1 if pretrained else None
    net = models.mobilenet_v2(weights=weights)
    net.classifier[1] = nn.Linear(net.last_channel, NUM_CLASSES)
    return net.to(DEVICE)


# ---------- 训练/验证 ----------
def train_one_epoch(model, loader, criterion, optimizer):
    model.train()
    running_loss, correct, total = 0.0, 0, 0
    for imgs, labels in tqdm(loader, desc='Train', leave=False):
        imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)
        optimizer.zero_grad()
        logits = model(imgs)
        loss = criterion(logits, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item() * imgs.size(0)

        # ★ 用 torch.eq 消除 bool 警告
        pred_mask: torch.Tensor = torch.eq(logits.argmax(1), labels)
        correct += int(pred_mask.sum())
        total += labels.size(0)
    return running_loss / total, correct / total


@torch.no_grad()
def evaluate(model, loader, criterion):
    model.eval()
    running_loss, correct, total = 0.0, 0, 0
    for imgs, labels in tqdm(loader, desc='Val', leave=False):
        imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)
        logits = model(imgs)
        loss = criterion(logits, labels)
        running_loss += loss.item() * imgs.size(0)

        # ★ 同上
        pred_mask: torch.Tensor = torch.eq(logits.argmax(1), labels)
        correct += int(pred_mask.sum())
        total += labels.size(0)
    return running_loss / total, correct / total


# ---------- 主流程 ----------
def main():
    criterion = nn.CrossEntropyLoss()
    if os.path.exists(MODEL_PATH):
        print('🚀 加载已有模型...')
        model = build_model(pretrained=False)
        model.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
    else:
        print('🔥 开始训练...')
        model = build_model(pretrained=True)
        optimizer = torch.optim.Adam(model.parameters(), lr=LR)
        best_acc, counter = 0.0, 0
        train_losses, train_accs, val_losses, val_accs = [], [], [], []

        for epoch in range(1, NUM_EPOCHS + 1):
            t0 = time.time()
            tr_loss, tr_acc = train_one_epoch(model, train_loader, criterion, optimizer)
            val_loss, val_acc = evaluate(model, val_loader, criterion)
            train_losses.append(tr_loss)
            train_accs.append(tr_acc)
            val_losses.append(val_loss)
            val_accs.append(val_acc)
            print(f'Epoch {epoch:02d}/{NUM_EPOCHS}  '
                  f'Train Loss:{tr_loss:.4f} Acc:{tr_acc:.4f}  '
                  f'Val Loss:{val_loss:.4f} Acc:{val_acc:.4f}  '
                  f'Time:{time.time() - t0:.1f}s')
            if val_acc > best_acc:
                best_acc, counter = val_acc, 0
                torch.save(model.state_dict(), MODEL_PATH)
            else:
                counter += 1
                if counter >= PATIENCE:
                    print('👉 Early stopping!')
                    break

        plt.figure(figsize=(12, 4))
        plt.subplot(1, 2, 1)
        plt.plot(train_losses, label='Train')
        plt.plot(val_losses, label='Val')
        plt.title('Loss')
        plt.legend()
        plt.grid()
        plt.subplot(1, 2, 2)
        plt.plot(train_accs, label='Train')
        plt.plot(val_accs, label='Val')
        plt.title('Accuracy')
        plt.legend()
        plt.grid()
        plt.tight_layout()
        plt.savefig(os.path.join(SAVE_DIR, 'curve.png'))
        plt.show()

    # ---------- 推理 ----------
    print('📝 预测集推理并生成带标注图片...')
    model.eval()
    from PIL import ImageDraw, ImageFont

    # 尽量用系统自带字体，避免路径问题
    try:
        font = ImageFont.truetype("arial.ttf", 18)
    except:
        font = ImageFont.load_default()

    pred_paths: List[str] = list(glob.glob(os.path.join(PRED_DIR, '*')))
    for path in tqdm(pred_paths, desc='Pred'):
        img = Image.open(path).convert('RGB')
        x = val_tf(img).unsqueeze(0).to(DEVICE)
        with torch.no_grad():
            prob = torch.softmax(model(x), dim=1).cpu().numpy().ravel()
        pred_cls = int(prob.argmax())

        # 在图片左上角画文字
        draw = ImageDraw.Draw(img)
        text = f"{train_set.classes[pred_cls]}  {prob[pred_cls]:.2f}"
        # 加黑色底色防止浅色图看不清
        draw.rectangle([(5, 5), (160, 25)], fill="black")
        draw.text((8, 5), text, fill="lime", font=font)

        # 保存到 printer，文件名与原图一致
        save_name = os.path.join(SAVE_DIR, "pred_" + os.path.basename(path))
        img.save(save_name)

    print('✅ 全部完成！带标注图片已保存到：', SAVE_DIR)


if __name__ == '__main__':
    main()